#!/usr/bin/python

import numpy,os,os.path,sys,time
from collections import Counter,defaultdict
from pylab import *
import tables,scipy
from scipy import fftpack
from scipy.cluster import vq
from scipy import stats
from scipy.ndimage import morphology,filters,interpolation
from ocrolib import docproc,improc
import pyflann
import random
from tables import openFile,Filters,Int32Atom,Float32Atom,Int64Atom
import multiprocessing
import fcntl
from llpy.sutils import *

import argparse
parser = argparse.ArgumentParser(description = "Concatenate data file.")
parser.add_argument('inputs',nargs="*")
parser.add_argument('-o','--output',help="output file")
parser.add_argument('-r','--pattern',default=None,help='pattern')
parser.add_argument('-b','--allowbad',action="store_true",help="allow bad classes")
parser.add_argument('-R','--norel',action='store_true',help="don't copy rel information")
parser.add_argument('-C','--nocosts',action='store_true',help="don't copy cost information")
parser.add_argument('-N','--nsamples',default=2000000000,type=int,help='approximate # output samples')
args = parser.parse_args()

with openFile(args.inputs[0]) as db:
    shape = list(db.root.patches[0].shape)
    print shape

select = None
if args.pattern is not None:
    select = re.compile('^'+args.pattern+'$')

with openFile(args.inputs[0]) as db:
    docosts = "costs" in dir(db.root) and not args.nocosts
    dorel = "rel" in dir(db.root) and not args.norel

if docosts: print "copying costs"
if dorel: print "copying rel"

with openFile(args.output,"w") as odb:
    table_log(odb,str(sys.argv))
    odb.createEArray(odb.root,"patches",tables.Float32Atom(),shape=[0]+shape,filters=tables.Filters(9))
    odb.createEArray(odb.root,"classes",tables.Int64Atom(),shape=(0,),filters=tables.Filters(9))
    if dorel:
        odb.createEArray(odb.root,"rel",tables.Float32Atom((3,)),shape=(0,),filters=tables.Filters(9))
    if docosts:
        odb.createEArray(odb.root,"costs",tables.Float32Atom(),shape=(0,),filters=tables.Filters(9))
    for fname in args.inputs:
        with openFile(fname) as db:
            table_lcopy(db,odb)
            nsamples = len(db.root.classes)
            for start,end in chunks(nsamples,10000):
                print "%9d %9d %s"%(len(odb.root.classes),start,fname)
                if len(odb.root.classes)>args.nsamples: break
                data = db.root.patches[start:end]
                clss = db.root.classes[start:end]
                if dorel: rel = db.root.rel[start:end]
                if docosts: costs = db.root.costs[start:end]
                if select is not None:
                    chars = [udecode(c) for c in clss]
                    good = array([(select.match(c) is not None) for c in chars],bool)
                    assert len(good)==len(data)
                    data = data[good]
                    clss = clss[good]
                    if dorel: rel = rel[good]
                    if docosts: costs = costs[good]
                # print [udecode(c) for c in clss[:10]]
                if not args.allowbad: assert (clss>32).all()
                n = min(end-start,args.nsamples-len(odb.root.patches))
                if n==0: break
                odb.root.patches.append(data[:n])
                odb.root.classes.append(clss[:n])
                if dorel: odb.root.rel.append(rel[:n])
                if docosts: odb.root.costs.append(costs[:n])
        if len(odb.root.classes)>args.nsamples: break
